#!/usr/local/bin/python2.7

from xml.etree.ElementTree import Element, ElementTree
import writer

class XMLWriter(writer.Writer):

    def __init__(self):
        writer.Writer.__init__(self)
        self._extension = 'xml'
        self._root = None

    def _writeIt(self):
        self._buildElement()

        tree = ElementTree(element = self._root)
        with open(self._dst, 'w') as myfile:
            tree.write(myfile)

        self._cleanXML()

    def getElement(self):
        self._buildElement()
        return self._root
    xml_element = property(getElement)

    def _buildElement(self):
        self._root = Element('component')
        dict_names = ['gen', 'mech', 'elec', 'iandc']

        for name in dict_names:
            curr_dict = self._component.getDict(name)
            self._dictToXML(name, curr_dict)

    def _dictToXML(self, cat_name, datamap):

        sub_cat = Element(cat_name)
        
        for key, vals in datamap.items():
            field = Element(key)

            for val in vals:
                field_val = Element('val')
                field_val.text = val
                field.append(field_val)
            
            sub_cat.append(field)

        self._root.append(sub_cat)

    def _cleanXML(self):
        """
        Takes an xml file and adds line breaks and (expanded) tabs to
        reflect the xml data hierarchy.

        This function is really a thin wrapper around the borrowed code of
        the *_indent* function defined below.
        """

        tree = ElementTree()
        tree.parse(self._dst) # load in the existing xml
        self._indent(tree.getroot()) # reformat it
        tree.write(self._dst, "utf-8") # rewrite it

    def _indent(self, elem, level = 0):
        """
        A recursive function that formats an xml.etree.ElementTree.Element
        object with line breaks and (expanded) tabs to reflect the xml data
        hierarchy.
        """

        tab_width = 2

        i = "\n" + level*tab_width*' '
        b = tab_width*' '
        if len(elem):
            if not elem.text or not elem.text.strip():
                elem.text = i + b
            for e in elem:
                self._indent(e, level + 1)
                if not e.tail or not e.tail.strip():
                    e.tail = i + b
            if not e.tail or not e.tail.strip():
                e.tail = i
        else:
            if level and (not elem.tail or not elem.tail.strip()):
                elem.tail = i

